import os
import pandas as pd
from PIL import Image
from PIL.ExifTags import TAGS, GPSTAGS
from PyQt5.QtWidgets import (QApplication, QWidget, QVBoxLayout, QHBoxLayout,
                             QPushButton, QFileDialog, QLabel, QSpinBox,
                             QTextEdit, QGroupBox, QMessageBox, QCheckBox, QRadioButton)
from PyQt5.QtGui import QTextCursor

from pypinyin import lazy_pinyin
import shutil


def chinese_sort_key(s):
    """中文文件名按拼音排序"""
    return lazy_pinyin(s)

class GPSInfoExtractor(QWidget):
    def __init__(self):
        super().__init__()
        self.flag = False
        self.init_ui()

    def init_ui(self):
        layout = QVBoxLayout()

        # 文件夹选择部分
        folder_group = QGroupBox("文件夹选择")
        folder_layout = QHBoxLayout()
        self.folder_label = QLabel("未选择文件夹")
        folder_btn = QPushButton("选择文件夹")
        folder_btn.clicked.connect(self.select_folder)
        folder_layout.addWidget(self.folder_label)
        folder_layout.addWidget(folder_btn)
        folder_group.setLayout(folder_layout)

        # 设置部分
        settings_group = QGroupBox("提取设置")
        settings_layout = QVBoxLayout()

        # 起始图片设置
        # start_layout = QHBoxLayout()
        # start_label = QLabel("从每个文件夹的第几张图片开始提取:")
        # self.start_spin = QSpinBox()
        # self.start_spin.setMinimum(1)
        # self.start_spin.setValue(1)
        # start_layout.addWidget(start_label)
        # start_layout.addWidget(self.start_spin)
        # start_layout.addStretch(1)

        self.radio_all = QRadioButton("提取所有图片")

        self.radio_start = QRadioButton("从第几张开始提取")
        self.spin_start = QSpinBox()
        self.spin_start.setMinimum(1)
        self.spin_start.setMaximum(99999)
        self.spin_start.setValue(1)
        self.radio_specific = QRadioButton("只提取第几张")
        self.radio_specific.setChecked(True)
        self.spin_specific = QSpinBox()
        self.spin_specific.setMinimum(1)
        self.spin_specific.setMaximum(99999)
        self.spin_specific.setValue(1)

        self.spin_last = QRadioButton("提取最后一张图片")

        option_layout1 = QHBoxLayout()
        option_layout1.addWidget(self.radio_all)

        option_layout2 = QHBoxLayout()
        option_layout2.addWidget(self.radio_start)
        option_layout2.addWidget(self.spin_start)

        option_layout3 = QHBoxLayout()
        option_layout3.addWidget(self.radio_specific)
        option_layout3.addWidget(self.spin_specific)

        option_layout4 = QHBoxLayout()
        option_layout4.addWidget(self.spin_last)



        # 输出格式选择
        format_layout = QHBoxLayout()
        self.excel_check = QCheckBox("Excel格式(.xlsx)")
        self.excel_check.setChecked(True)
        self.csv_check = QCheckBox("CSV格式(.csv)")
        self.csv_check.setChecked(True)
        format_layout.addWidget(QLabel("输出格式:"))
        format_layout.addWidget(self.excel_check)
        format_layout.addWidget(self.csv_check)
        format_layout.addStretch(1)

        # settings_layout.addLayout(start_layout)
        settings_layout.addLayout(option_layout3)
        settings_layout.addLayout(option_layout2)
        settings_layout.addLayout(option_layout4)
        settings_layout.addLayout(option_layout1)




        settings_layout.addLayout(format_layout)
        settings_group.setLayout(settings_layout)

        # 操作按钮
        btn_group = QGroupBox("操作")
        btn_layout = QHBoxLayout()
        self.extract_btn = QPushButton("开始提取")
        self.extract_btn.clicked.connect(self.extract_gps_info)

        self.extract_stop_btn = QPushButton("停止提取")
        self.extract_stop_btn.clicked.connect(self.extract_stop_gps_info)

        clear_btn = QPushButton("清空输出")
        clear_btn.clicked.connect(self.clear_output)
        btn_layout.addWidget(self.extract_btn)
        btn_layout.addWidget(self.extract_stop_btn)
        btn_layout.addWidget(clear_btn)
        btn_group.setLayout(btn_layout)

        # 输出区域
        output_group = QGroupBox("提取进度")
        output_layout = QVBoxLayout()
        self.output_text = QTextEdit()
        self.output_text.setReadOnly(True)
        output_layout.addWidget(self.output_text)
        output_group.setLayout(output_layout)

        # 主布局
        layout.addWidget(folder_group)
        layout.addWidget(settings_group)
        layout.addWidget(btn_group)
        layout.addWidget(output_group)

        self.setLayout(layout)
        self.setWindowTitle("GPS信息提取工具")
        self.resize(800, 600)

    def select_folder(self):
        folder = QFileDialog.getExistingDirectory(self, "选择文件夹")
        if folder:
            self.folder_label.setText(folder)
            self.append_output(f"已选择文件夹: {folder}")


    def extract_stop_gps_info(self):
        self.flag = True

    def extract_gps_info(self):
        folder = self.folder_label.text()
        if folder == "未选择文件夹":
            self.append_output("错误: 请先选择文件夹", error=True)
            return

        if not self.excel_check.isChecked() and not self.csv_check.isChecked():
            self.append_output("错误: 请至少选择一种输出格式", error=True)
            return

        # start_index = self.start_spin.value() - 1  # 转换为0-based

        # 确定提取模式
        if self.radio_all.isChecked():
            mode = "all"
        elif self.radio_start.isChecked():
            mode = "start"
        elif self.spin_last.isChecked():
            mode = "last"
        else:
            mode = "specific"



        # 准备数据列表
        data = []
        data_csv = []
        total_images = 0
        processed_images = 0
        has_gps_count = 0
        self.extract_btn.setEnabled(False)

        # 创建图片复制文件夹
        image_copy_dir = os.path.join(os.path.dirname(folder), "GPS信息图片")
        os.makedirs(image_copy_dir, exist_ok=True)


        # 先统计总图片数
        self.append_output("正在扫描文件夹中的图片...")
        QApplication.processEvents()

        for root, dirs, files in os.walk(folder):
            for file in files:
                if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                    total_images += 1

        if total_images == 0:
            self.append_output("错误: 文件夹中没有找到图片文件", error=True)
            return

        self.append_output(f"共找到 {total_images} 张图片，开始提取GPS信息...")


        # 遍历所有文件夹

        for root, dirs, files in os.walk(folder):
            if self.flag is True:
                self.extract_btn.setEnabled(True)
                self.flag = False
                return

            dirs.sort(key=chinese_sort_key)  # 影响子目录遍历顺序
            files.sort(key=chinese_sort_key)  # 当前目录文件排序
            # 获取当前文件夹中的图片文件
            image_files = []

            for file in files:
                if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                    image_files.append(os.path.join(root, file))



            if not image_files:
                continue

            # 只处理从start_index开始的图片
            if mode == 'all':
                start_index = 0
                current_folder_images = image_files[start_index:] if start_index < len(image_files) else []
            elif mode == 'start':
                start_index = self.spin_start.value() - 1
                current_folder_images = image_files[start_index:] if start_index < len(image_files) else []
            elif mode == 'last':
                start_index = len(image_files) - 1
                current_folder_images = image_files[start_index:] if start_index < len(image_files) else []

            else:
                specific_num = self.spin_specific.value() - 1
                current_folder_images = [image_files[specific_num]] if specific_num < len(image_files) else []


            for img_path in current_folder_images:

                try:
                    processed_images += 1
                    progress = f"[{processed_images}/{total_images}]"
                    self.append_output(f"{progress} 正在处理: {os.path.relpath(img_path, folder)}")
                    QApplication.processEvents()

                    gps_info = self.get_gps_info(img_path)
                    if gps_info:
                        has_gps_count += 1
                        row = {
                            '图片名称': os.path.basename(os.path.basename(img_path)),
                            '图片所在文件夹': os.path.basename(os.path.dirname(img_path)),
                            '纬度': gps_info.get('Latitude', '无'),
                            '经度': gps_info.get('Longitude', '无'),
                            '海拔': gps_info.get('GPSAltitude', '无'),
                            '拍摄时间': gps_info.get('GPSDateStamp', '无'),
                        }
                        data.append(row)

                        row_csv = {
                            '图片所在文件夹': os.path.basename(os.path.dirname(img_path)),
                            '纬度': gps_info.get('Latitude', '无'),
                            '经度': gps_info.get('Longitude', '无'),
                        }
                        data_csv.append(row_csv)

                        self.append_output(f"  找到GPS信息: 纬度 {row['纬度']}, 经度 {row['经度']}")
                    else:
                        row = {
                            '图片名称': os.path.basename(os.path.basename(img_path)),
                            '图片所在文件夹': os.path.basename(os.path.dirname(img_path)),
                            '纬度': '无',
                            '经度': '无',
                            '海拔': '无',
                            '拍摄时间': '无',
                        }
                        data.append(row)

                        row_csv = {
                            '图片所在文件夹': os.path.basename(os.path.dirname(img_path)),
                            '纬度': '无',
                            '经度': '无',
                        }
                        data_csv.append(row_csv)

                        self.append_output("  无GPS信息")
                except Exception as e:
                    self.extract_btn.setEnabled(True)
                    self.append_output(f"  处理图片出错: {str(e)}", error=True)

            # 复制图片
            copied_images = 0
            for img_path in current_folder_images:
                # 复制图片并重命名为目录名（直接覆盖同名文件）
                try:
                    # 获取图片所在目录的名称
                    dir_name = os.path.basename(os.path.dirname(img_path))
                    # 获取文件扩展名
                    ext = os.path.splitext(img_path)[1].lower()

                    # 构造新文件名
                    new_name = f"{dir_name}{ext}"
                    dest_path = os.path.join(image_copy_dir, new_name)

                    # 直接复制并覆盖同名文件
                    shutil.copy2(img_path, dest_path)
                    copied_images += 1
                    self.append_output(f"已复制图片并重命名为: {new_name}")
                except Exception as e:
                    self.append_output(f"复制图片 {img_path} 失败: {str(e)}")


        self.extract_btn.setEnabled(True)


        # 保存文件
        if data:

            current_dir = os.path.dirname(folder)
            output_folder = os.path.join(current_dir, "坐标")
            os.makedirs(output_folder, exist_ok=True)

            base_path = os.path.join(output_folder, "GPS")
            success_formats = []

            if self.excel_check.isChecked():
                excel_path = f"{base_path}.xlsx"
                try:
                    df = pd.DataFrame(data)
                    df.to_excel(excel_path, index=False, engine='openpyxl')
                    success_formats.append(f"Excel文件({excel_path})")
                    self.append_output(f"Excel文件已保存: {excel_path}")
                except Exception as e:
                    self.append_output(f"错误: 保存Excel文件时出错 - {str(e)}", error=True)

            if self.csv_check.isChecked():
                csv_path = f"{base_path}.csv"
                try:
                    df = pd.DataFrame(data_csv)
                    df.to_csv(csv_path, index=False, encoding='utf_8_sig')  # 使用utf_8_sig编码支持中文
                    success_formats.append(f"CSV文件({csv_path})")
                    self.append_output(f"CSV文件已保存: {csv_path}")
                except Exception as e:
                    self.append_output(f"错误: 保存CSV文件时出错 - {str(e)}", error=True)

            if success_formats:
                summary_msg = (
                    f"完成! 共处理 {processed_images} 张图片\n"
                    f"其中 {has_gps_count} 张包含GPS信息\n"
                    f"已保存为: {', '.join(success_formats)}"
                )
                self.append_output(summary_msg)
                QMessageBox.information(self, "提取完成", summary_msg)
        else:
            self.append_output("错误: 没有提取到任何GPS信息", error=True)

    def get_gps_info(self, image_path):
        """从图片中提取GPS信息"""
        try:
            img = Image.open(image_path)
            exif_data = img._getexif()

            if not exif_data:
                return None

            # 获取常规EXIF信息
            exif = {}
            for tag, value in exif_data.items():
                decoded = TAGS.get(tag, tag)
                exif[decoded] = value

            # 获取GPS信息
            gps_info = {}
            if 'GPSInfo' in exif:
                for key in exif['GPSInfo'].keys():
                    decode = GPSTAGS.get(key, key)
                    gps_info[decode] = exif['GPSInfo'][key]

                # 转换经纬度为十进制
                lat = self.convert_to_degrees(gps_info.get('GPSLatitude'))
                if 'GPSLatitudeRef' in gps_info and gps_info['GPSLatitudeRef'] == 'S':
                    lat = -lat

                lon = self.convert_to_degrees(gps_info.get('GPSLongitude'))
                if 'GPSLongitudeRef' in gps_info and gps_info['GPSLongitudeRef'] == 'W':
                    lon = -lon

                # 海拔
                altitude = gps_info.get('GPSAltitude')
                if altitude:
                    altitude = float(altitude)


                return {
                    "Latitude": lat,
                    "Longitude": lon,
                    "GPSAltitude": altitude,
                    'GPSLatitude': gps_info.get('GPSLatitude'),
                    'GPSLongitude': gps_info.get('GPSLongitude'),
                    'DateTimeOriginal': gps_info.get('DateTimeOriginal'),
                    'GPSTimeStamp': gps_info.get('GPSTimeStamp'),
                    "GPSDateStamp": gps_info.get('GPSDateStamp'),
                    'GPSStatus': gps_info.get('GPSStatus'),
                    'GPSMeasureMode': gps_info.get('GPSMeasureMode')
                }
            return None
        except Exception as e:
            self.output_text.append(f"提取 {os.path.basename(image_path)} 的GPS信息时出错: {str(e)}")
            return None

    def convert_to_degrees(self, value):
        try:
            if isinstance(value, tuple) and len(value) == 3:
                d, m, s = value
                return float(d) + (float(m) / 60.0) + (float(s) / 3600.0)
            return 0.0
        except:
            return 0.0

    def append_output(self, text, error=False):
        cursor = self.output_text.textCursor()
        cursor.movePosition(QTextCursor.End)
        if error:
            cursor.insertHtml(f'<font color="red">{text}</font><br>')
        else:
            cursor.insertHtml(f'{text}<br>')
        self.output_text.setTextCursor(cursor)
        self.output_text.ensureCursorVisible()

    def clear_output(self):
        self.output_text.clear()

